﻿///
/// @file AST.cpp
/// @brief 抽象语法树AST管理的实现
/// @author zenglj (zenglj@live.com)
/// @version 1.1
/// @date 2024-11-23
///
/// @copyright Copyright (c) 2024
///
/// @par 修改日志:
/// <table>
/// <tr><th>Date       <th>Version <th>Author  <th>Description
/// <tr><td>2024-11-21 <td>1.0     <td>zenglj  <td>新做
/// <tr><td>2024-11-23 <td>1.1     <td>zenglj  <td>表达式版增强
/// </table>
///
#include <cstdarg>
#include <cstdint>
#include <string>

#include "AST.h"
#include "AttrType.h"
#include "Types/IntegerType.h"
#include "Types/VoidType.h"
#include "Types/FloatType.h"
#include <iostream>

/* 整个AST的根节点 */
ast_node * ast_root = nullptr;

/// @brief 创建指定节点类型的节点
/// @param _node_type 节点类型
/// @param _line_no 行号
ast_node::ast_node(ast_operator_type _node_type, Type * _type, int64_t _line_no)
    : node_type(_node_type), line_no(-1), type(_type)
{}

/// @brief 构造函数
/// @param _type 节点值的类型
/// @param line_no 行号
ast_node::ast_node(Type * _type) : ast_node(ast_operator_type::AST_OP_LEAF_TYPE, _type, -1)
{}

/// @brief 针对无符号整数字面量的构造函数
/// @param attr 无符号整数字面量
ast_node::ast_node(digit_int_attr attr)
    : ast_node(ast_operator_type::AST_OP_LEAF_LITERAL_UINT, IntegerType::getTypeInt(), attr.lineno)
{
    integer_val = attr.val;
}

/// @brief 针对浮点数字面量的构造函数
/// @param _float_val 浮点数字面量
/// @param _line_no 行号
ast_node::ast_node(float _float_val, int64_t _line_no)
    : ast_node(ast_operator_type::AST_OP_LEAF_LITERAL_FLOAT, FloatType::getTypeFloat(), _line_no)
{
    float_val = _float_val;
}

/// @brief 针对标识符ID的叶子构造函数
/// @param attr 字符型字面量
ast_node::ast_node(var_id_attr attr) : ast_node(ast_operator_type::AST_OP_LEAF_VAR_ID, VoidType::getType(), attr.lineno)
{
    name = attr.id;
}

/// @brief 针对标识符ID的叶子构造函数
/// @param _id 标识符ID
/// @param _line_no 行号
ast_node::ast_node(std::string _id, int64_t _line_no)
    : ast_node(ast_operator_type::AST_OP_LEAF_VAR_ID, VoidType::getType(), _line_no)
{
    name = _id;
}

/// @brief 判断是否是叶子节点
/// @return true：是叶子节点 false：内部节点
bool ast_node::isLeafNode()
{
    bool is_leaf;

    switch (this->node_type) {
        case ast_operator_type::AST_OP_LEAF_LITERAL_UINT:
        case ast_operator_type::AST_OP_LEAF_LITERAL_FLOAT:
        case ast_operator_type::AST_OP_LEAF_VAR_ID:
        case ast_operator_type::AST_OP_LEAF_TYPE:
        // 增加形参相关节点判断，视作叶子节点
        case ast_operator_type::AST_OP_FUNC_FORMAL_PARAM_TYPE:
        case ast_operator_type::AST_OP_FUNC_FORMAL_PARAM_NAME:
        case ast_operator_type::AST_OP_FUNC_FORMAL_PARAM_DIM:
            is_leaf = true;
            break;
        default:
            is_leaf = false;
            break;
    }

    return is_leaf;
}

/// @brief 创建指定节点类型的节点，请注意在指定有效的孩子后必须追加一个空指针nullptr，表明可变参数结束
/// @param type 节点类型
/// @param son_num 孩子节点的个数
/// @param ...
/// 可变参数，可支持插入若干个孩子节点，自左往右的次序，最後一个孩子节点必须指定为nullptr。如果没有孩子，则指定为nullptr
/// @return 创建的节点
ast_node * ast_node::New(ast_operator_type type, ...)
{
    ast_node * parent_node = new ast_node(type);

    va_list valist;

    /* valist指向传入的第一个可选参数 */
    va_start(valist, type);

    for (;;) {

        // 获取节点对象。如果最後一个对象为空指针，则说明结束
        ast_node * node = va_arg(valist, ast_node *);
        if (nullptr == node) {
            break;
        }

        // 插入到父节点中
        parent_node->insert_son_node(node);
    }

    /* 清理为 valist 保留的内存 */
    va_end(valist);

    return parent_node;
}

/// @brief 向父节点插入一个节点
/// @param parent 父节点
/// @param node 节点
ast_node * ast_node::insert_son_node(ast_node * node)
{
    if (node) {

        // 孩子节点有效时加入，主要为了避免空语句等时会返回空指针
        node->parent = this;
        this->sons.push_back(node);
    }

    return this;
}

/// @brief 创建无符号整数的叶子节点
/// @param attr 无符号整数字面量
ast_node * ast_node::New(digit_int_attr attr)
{
    ast_node * node = new ast_node(attr);

    return node;
}

/// @brief 创建标识符的叶子节点
/// @param attr 字符型字面量
ast_node * ast_node::New(var_id_attr attr)
{
    ast_node * node = new ast_node(attr);

    return node;
}

/// @brief 创建标识符的叶子节点
/// @param id 词法值
/// @param line_no 行号
ast_node * ast_node::New(std::string id, int64_t lineno)
{
    ast_node * node = new ast_node(id, lineno);

    return node;
}

/// @brief 创建具备指定类型的节点
/// @param type 节点值类型
/// @param line_no 行号
/// @return 创建的节点
ast_node * ast_node::New(Type * type)
{
    ast_node * node = new ast_node(type);

    return node;
}

/// @brief 递归清理抽象语法树
/// @param node AST的节点
void ast_node::Delete(ast_node * node)
{
    if (node) {

        for (auto child: node->sons) {
            ast_node::Delete(child);
        }

        //// 这里没有必要清理孩子，由于下面就要删除该节点
        node->sons.clear();
    }

    // 清理node资源
    // delete node;
}

///
/// @brief AST资源清理
///
void free_ast(ast_node * root)
{
    ast_node::Delete(root);
}

/// @brief 创建函数定义类型的内部AST节点
/// @param type_node 类型节点
/// @param name_node 函数名字节点
/// @param block_node 函数体语句块节点
/// @param params_node 函数形参，可以没有参数
/// @return 创建的节点
ast_node * create_func_def(ast_node * type_node, ast_node * name_node, ast_node * block_node, ast_node * params_node)
{
    ast_node * node = new ast_node(ast_operator_type::AST_OP_FUNC_DEF, type_node->type, name_node->line_no);

    // 设置函数名
    node->name = name_node->name;

    // 如果没有参数，则创建参数节点
    if (!params_node) {
        params_node = new ast_node(ast_operator_type::AST_OP_FUNC_FORMAL_PARAMS);
    }

    // 如果没有函数体，则创建函数体，也就是语句块
    if (!block_node) {
        block_node = new ast_node(ast_operator_type::AST_OP_BLOCK);
    }

    (void) node->insert_son_node(type_node);
    (void) node->insert_son_node(name_node);
    (void) node->insert_son_node(params_node);
    (void) node->insert_son_node(block_node);

    return node;
}

/// @brief 创建函数定义类型的内部AST节点
/// @param type 返回值类型
/// @param id 函数名字
/// @param block_node 函数体语句块节点
/// @param params_node 函数形参，可以没有参数
/// @return 创建的节点
ast_node * create_func_def(type_attr & type, var_id_attr & id, ast_node * block_node, ast_node * params_node)
{
    // 创建整型类型节点的终结符节点
    ast_node * type_node = create_type_node(type);

    // 创建标识符终结符节点
    ast_node * id_node = ast_node::New(id.id, id.lineno);

    // 对于字符型字面量的字符串空间需要释放，因词法用到了strdup进行了字符串复制
    free(id.id);
    id.id = nullptr;

    return create_func_def(type_node, id_node, block_node, params_node);
}

/// @brief 创建AST的内部节点
/// @param node_type 节点类型
/// @param first_child 第一个孩子节点
/// @param second_child 第一个孩子节点
/// @param third_child 第一个孩子节点
/// @return 创建的节点
ast_node * create_contain_node(ast_operator_type node_type,
                               ast_node * first_child,
                               ast_node * second_child,
                               ast_node * third_child)
{
    ast_node * node = new ast_node(node_type);

    if (first_child) {
        (void) node->insert_son_node(first_child);
    }

    if (second_child) {
        (void) node->insert_son_node(second_child);
    }

    if (third_child) {
        (void) node->insert_son_node(third_child);
    }

    return node;
}

Type * typeAttr2Type(type_attr & attr)
{
    if (attr.type == BasicType::TYPE_INT) {
        return IntegerType::getTypeInt();
    } else if (attr.type == BasicType::TYPE_FLOAT) {
        return FloatType::getTypeFloat();
    } else {
        return VoidType::getType();
    }
}

/// @brief 创建类型节点
/// @param type 类型信息
/// @return 创建的节点
ast_node * create_type_node(type_attr & attr)
{
    Type * type = typeAttr2Type(attr);

    ast_node * type_node = ast_node::New(type);

    return type_node;
}
// Type * typeAttr2Type(type_attr & attr)
// {
//     Type * t = nullptr;
//     if (attr.type == BasicType::TYPE_INT) {
//         t = IntegerType::getTypeInt();
//     } else if (attr.type == BasicType::TYPE_FLOAT) {
//         t = FloatType::getTypeFloat();
//     } else {
//         t = VoidType::getType();
//     }
//     // 调试输出
//     std::cout << "[DEBUG] typeAttr2Type: attr.type=" << static_cast<int>(attr.type) << " -> "
//               << (t ? t->toString() : "null") << std::endl;
//     return t;
// }

// ast_node * create_type_node(type_attr & attr)
// {
//     Type * type = typeAttr2Type(attr);

//     // 调试输出
//     std::cout << "[DEBUG] create_type_node: attr.type=" << static_cast<int>(attr.type)
//               << " type=" << (type ? type->toString() : "null") << std::endl;

//     ast_node * type_node = ast_node::New(type);

//     return type_node;
// }

/// @brief 创建函数调用的节点
/// @param funcname_node 函数名节点
/// @param params_node 实参节点
/// @return 创建的节点
ast_node * create_func_call(ast_node * funcname_node, ast_node * params_node)
{
    ast_node * node = new ast_node(ast_operator_type::AST_OP_FUNC_CALL);

    // 设置调用函数名
    if (funcname_node->name == "starttime" || funcname_node->name == "stoptime") {
        funcname_node->name = "_sysy_" + funcname_node->name; // 特殊处理starttime和stoptime
    }
    node->name = funcname_node->name;

    // 如果没有参数，则创建参数节点
    if (!params_node) {
        params_node = new ast_node(ast_operator_type::AST_OP_FUNC_REAL_PARAMS);
    }

    (void) node->insert_son_node(funcname_node);
    (void) node->insert_son_node(params_node);

    return node;
}

///
/// @brief 根据第一个变量定义创建变量声明语句节点
/// @param first_child 第一个变量定义节点，其类型为AST_OP_VAR_DECL
/// @return ast_node* 变量声明语句节点
///
ast_node * create_var_decl_stmt_node(ast_node * first_child)
{
    // 创建变量声明语句
    ast_node * stmt_node = create_contain_node(ast_operator_type::AST_OP_DECL_STMT);

    if (first_child) {

        stmt_node->type = first_child->type;

        // 插入到变量声明语句
        (void) stmt_node->insert_son_node(first_child);
    }

    return stmt_node;
}

ast_node * createVarDeclNode(Type * type, var_id_attr & id)
{
    // 创建整型类型节点的终结符节点
    ast_node * type_node = ast_node::New(type);

    // 创建标识符终结符节点
    ast_node * id_node = ast_node::New(id.id, id.lineno);

    // 对于字符型字面量的字符串空间需要释放，因词法用到了strdup进行了字符串复制
    free(id.id);
    id.id = nullptr;

    // 创建变量定义节点
    ast_node * decl_node = create_contain_node(ast_operator_type::AST_OP_VAR_DECL, type_node, id_node);

    // 暂存类型
    decl_node->type = type;

    return decl_node;
}

ast_node * createVarDeclNode(type_attr & type, var_id_attr & id)
{
    return createVarDeclNode(typeAttr2Type(type), id);
}

///
/// @brief 根据变量的类型和属性创建变量声明语句节点
/// @param type 变量的类型
/// @param id 变量的名字
/// @return ast_node* 变量声明语句节点
///
ast_node * create_var_decl_stmt_node(type_attr & type, var_id_attr & id)
{
    // 创建变量定义节点
    ast_node * decl_node = createVarDeclNode(type, id);

    // 创建变量声明语句
    ast_node * stmt_node = create_contain_node(ast_operator_type::AST_OP_DECL_STMT);

    stmt_node->type = decl_node->type;

    // 插入到变量声明语句
    (void) stmt_node->insert_son_node(decl_node);

    return stmt_node;
}

///
/// @brief 向变量声明语句中追加变量声明
/// @param stmt_node 变量声明语句
/// @param id 变量的名字
/// @return ast_node* 变量声明语句节点
///
ast_node * add_var_decl_node(ast_node * stmt_node, var_id_attr & id)
{
    // 创建变量定义节点
    ast_node * decl_node = createVarDeclNode(stmt_node->type, id);

    // 插入到变量声明语句
    (void) stmt_node->insert_son_node(decl_node);

    return stmt_node;
}

// 修改：创建函数形式参数结点的实现
ast_node * create_func_formal_param(uint32_t line_no, const char * param_name, Type * param_type, ast_node * dim_node)
{
    ast_node * param = new ast_node(ast_operator_type::AST_OP_FUNC_FORMAL_PARAM);
    param->line_no = line_no;
    param->type = param_type;

    // 创建类型叶子节点
    ast_node * typeNode = new ast_node(ast_operator_type::AST_OP_FUNC_FORMAL_PARAM_TYPE);
    typeNode->type = param_type;
    typeNode->name = param_type->toString();
    param->insert_son_node(typeNode);

    // 创建名称叶子节点
    ast_node * nameNode = new ast_node(ast_operator_type::AST_OP_FUNC_FORMAL_PARAM_NAME);
    nameNode->name = param_name;
    param->insert_son_node(nameNode);

    // 添加维度节点（如果存在）
    if (dim_node) {
        param->insert_son_node(dim_node);
    }

    return param;
}

/// @brief 创建数组访问的节点
/// @param array_node 数组节点
/// @param dim_node 维度节点
/// @return 创建的节点
ast_node * create_array_access(ast_node * array_node, ast_node * dim_node)
{
    ast_node * node = new ast_node(ast_operator_type::AST_OP_ARRAY_ACCESS);

    // 设置数组名
    node->name = array_node->name;

    // 添加数组标识符和维度作为子节点
    (void) node->insert_son_node(array_node);
    (void) node->insert_son_node(dim_node);

    return node;
}

// /// @brief 创建宏定义声明节点
// /// @param name_node 宏名称节点
// /// @param value_node 宏值表达式节点
// /// @return 创建的节点
// ast_node * create_define_decl(ast_node * name_node, ast_node * value_node)
// {
//     ast_node * declNode = new ast_node(ast_operator_type::AST_OP_DEFINE_DECL);

//     // 创建宏定义节点
//     ast_node * defNode = new ast_node(ast_operator_type::AST_OP_DEFINE_DEF);
//     defNode->name = name_node->name;
//     defNode->line_no = name_node->line_no;

//     defNode->insert_son_node(name_node);
//     defNode->insert_son_node(value_node);

//     declNode->insert_son_node(defNode);

//     return declNode;
// }

// /// @brief 创建数组维度节点
// /// @param dim_expr 维度表达式节点
// /// @return 创建的节点
// ast_node * create_array_dim(ast_node * dim_expr)
// {
//     ast_node * node = new ast_node(ast_operator_type::AST_OP_ARRAY_DIM);

//     // 添加维度表达式作为子节点
//     (void) node->insert_son_node(dim_expr);

//     return node;
// }

// /// @brief 创建类型转换节点
// /// @param type_node 目标类型节点
// /// @param expr_node 被转换的表达式节点
// /// @return 创建的节点
// ast_node * create_cast_node(ast_node * type_node, ast_node * expr_node)
// {
//     ast_node * node = new ast_node(ast_operator_type::AST_OP_CAST);

//     // 添加类型和表达式作为子节点
//     (void) node->insert_son_node(type_node);
//     (void) node->insert_son_node(expr_node);

//     return node;
// }

// /// @brief 创建指针类型节点
// /// @param base_type 基础类型
// /// @param ptr_level 指针级别
// /// @return 创建的节点
// ast_node * create_pointer_type_node(Type * base_type, int ptr_level)
// {
//     ast_node * node = new ast_node(ast_operator_type::AST_OP_POINTER_TYPE);

//     node->type = base_type;
//     node->integer_val = ptr_level; // 使用integer_val存储指针级别

//     std::string typeName = base_type->toString();
//     for (int i = 0; i < ptr_level; i++) {
//         typeName += "*";
//     }
//     node->name = typeName;

//     return node;
// }

// /// @brief 根据基本类型名称创建指针类型节点
// /// @param type_name 类型名称 ("int", "float", "void")
// /// @param ptr_level 指针级别
// /// @return 创建的节点
// ast_node * create_pointer_type_node(const std::string & type_name, int ptr_level)
// {
//     Type * base_type;

//     if (type_name == "int") {
//         base_type = IntegerType::getTypeInt();
//     } else if (type_name == "float") {
//         base_type = FloatType::getTypeFloat();
//     } else if (type_name == "void") {
//         base_type = VoidType::getType();
//     } else {
//         // 默认为void类型
//         base_type = VoidType::getType();
//     }

//     return create_pointer_type_node(base_type, ptr_level);
// }

/// @brief 检查节点是否可以进行单目运算符优化
/// @param node AST节点
/// @return true: 可以优化 false: 不可优化
bool canOptimizeUnary(ast_node * node)
{
    if (!node)
        return false;

    return (node->node_type == ast_operator_type::AST_OP_UNARY_PLUS ||
            node->node_type == ast_operator_type::AST_OP_UNARY_MINUS ||
            node->node_type == ast_operator_type::AST_OP_UNARY_NOT);
}

/// @brief 收集所有连续的单目运算符序列
/// @param node AST节点
/// @param operators 输出参数，存储运算符序列（运算符类型和行号）
/// @return 最内层的操作数节点
ast_node * collectUnaryOperators(ast_node * node, std::vector<std::pair<ast_operator_type, int64_t>> & operators)
{
    if (!node || !canOptimizeUnary(node)) {
        return node;
    }

    // 存储运算符类型和行号
    operators.push_back(std::make_pair(node->node_type, node->line_no));

    if (node->sons.empty()) {
        return nullptr;
    }

    return collectUnaryOperators(node->sons[0], operators);
}

/// @brief 优化单目运算符表达式
/// @param node 待优化的AST节点
/// @return 优化后的节点
ast_node * optimizeUnaryExpr(ast_node * node)
{
    if (!node || !canOptimizeUnary(node)) {
        return node;
    }

    // 收集所有连续的单目运算符（外到内）
    std::vector<std::pair<ast_operator_type, int64_t>> operators;
    ast_node * innerOperand = collectUnaryOperators(node, operators);

    if (!innerOperand) {
        return node;
    }

    // 递归优化操作数
    innerOperand = optimizeAST(innerOperand);

    // 逆序（内到外）处理，保持C代码顺序
    ast_node * result = innerOperand;
    for (int i = (int) operators.size() - 1; i >= 0; --i) {
        ast_operator_type op = operators[i].first;
        int64_t line_no = operators[i].second;

        // 连续的+可以直接跳过
        if (op == ast_operator_type::AST_OP_UNARY_PLUS) {
            continue;
        }

        // 连续的-和!需要折叠：如果前面有同类型的连续偶数个，可以抵消
        // 但要保持顺序，所以只在遇到连续偶数个时跳过
        // 这里只对!和-做折叠
        if (op == ast_operator_type::AST_OP_UNARY_MINUS) {
            // 检查有多少连续-
            int cnt = 1;
            while (i - 1 >= 0 && operators[i - 1].first == ast_operator_type::AST_OP_UNARY_MINUS) {
                ++cnt;
                --i;
            }
            if (cnt % 2 == 1) {
                ast_node * negNode = new ast_node(ast_operator_type::AST_OP_UNARY_MINUS, nullptr, line_no);
                negNode->type = result->type;
                negNode->insert_son_node(result);
                result = negNode;
            }
            continue;
        }
        if (op == ast_operator_type::AST_OP_UNARY_NOT) {
            // 检查有多少连续!
            int cnt = 1;
            while (i - 1 >= 0 && operators[i - 1].first == ast_operator_type::AST_OP_UNARY_NOT) {
                ++cnt;
                --i;
            }
            if (cnt % 2 == 1) {
                ast_node * notNode = new ast_node(ast_operator_type::AST_OP_UNARY_NOT, nullptr, line_no);
                notNode->type = result->type;
                notNode->insert_son_node(result);
                result = notNode;
            }
            continue;
        }
    }

    return result;
}

/// @brief 检查节点及其子树是否只包含单目运算符
/// @param node AST节点
/// @return true: 只包含单目运算符 false: 包含其他类型节点
bool isUnaryChain(ast_node * node)
{
    if (!node)
        return false;

    if (!canOptimizeUnary(node)) {
        return false;
    }

    if (node->sons.empty()) {
        return false; // 没有操作数
    }

    if (node->sons.size() > 1) {
        return false; // 单目运算符只能有一个操作数
    }

    ast_node * child = node->sons[0];

    // 如果子节点也是单目运算符，继续检查
    if (canOptimizeUnary(child)) {
        return isUnaryChain(child);
    }

    // 子节点不是单目运算符，说明是操作数，这是一个有效的单目运算符链
    return true;
}

/// @brief 递归优化整个AST树
/// @param root AST根节点
/// @return 优化后的AST根节点
ast_node * optimizeAST(ast_node * root)
{
    if (!root) {
        return nullptr;
    }

    // 先递归优化子节点
    for (size_t i = 0; i < root->sons.size(); i++) {
        root->sons[i] = optimizeAST(root->sons[i]);
    }

    // 检查当前节点是否是单目运算符链的开始
    if (canOptimizeUnary(root) && isUnaryChain(root)) {
        return optimizeUnaryExpr(root);
    }

    return root;
}